#!/usr/bin/env python
"""
百家乐系统数据库管理工具
提供数据库管理功能：查看表信息，清空表，导出数据等
"""
import os
import sys
import json
import argparse
import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 从init_db导入函数
from init_db import init_database, check_database_exists

def get_db_connection():
    """获取数据库连接"""
    load_dotenv()
    
    # 数据库连接参数
    db_config = {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 5432)),
        'user': os.getenv('DB_USER', 'postgres'),
        'password': os.getenv('DB_PASSWORD', 'Postgre,.1'),
        'database': os.getenv('DB_NAME', 'baijiale')
    }
    
    try:
        conn = psycopg2.connect(**db_config)
        return conn
    except Exception as e:
        print(f"数据库连接错误: {str(e)}")
        return None

def list_tables():
    """列出所有表"""
    conn = get_db_connection()
    if not conn:
        return False
    
    cursor = None
    try:
        cursor = conn.cursor()
        cursor.execute("""
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = 'public'
            ORDER BY table_name
        """)
        
        tables = cursor.fetchall()
        
        print("\n数据库中的表:")
        print("-" * 30)
        
        if not tables:
            print("数据库中没有表")
        else:
            for i, table in enumerate(tables, 1):
                print(f"{i}. {table[0]}")
        
        return True
    
    except Exception as e:
        print(f"列出表错误: {str(e)}")
        return False
    
    finally:
        if cursor:
            cursor.close()
        conn.close()

def describe_table(table_name):
    """描述表结构"""
    conn = get_db_connection()
    if not conn:
        return False
    
    cursor = None
    try:
        cursor = conn.cursor()
        
        # 检查表是否存在
        cursor.execute("""
            SELECT 1 
            FROM information_schema.tables 
            WHERE table_schema = 'public' AND table_name = %s
        """, (table_name,))
        
        if not cursor.fetchone():
            print(f"表 '{table_name}' 不存在")
            return False
        
        # 获取表结构
        cursor.execute("""
            SELECT column_name, data_type, is_nullable, column_default
            FROM information_schema.columns
            WHERE table_schema = 'public' AND table_name = %s
            ORDER BY ordinal_position
        """, (table_name,))
        
        columns = cursor.fetchall()
        
        # 获取表行数
        cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
        row_count = cursor.fetchone()[0]
        
        print(f"\n表 '{table_name}' 结构 (共 {row_count} 行):")
        print("-" * 80)
        print(f"{'列名':<20} {'数据类型':<15} {'可空':<6} {'默认值':<20}")
        print("-" * 80)
        
        for col in columns:
            col_name, data_type, is_nullable, default = col
            print(f"{col_name:<20} {data_type:<15} {is_nullable:<6} {str(default or ''):<20}")
        
        return True
    
    except Exception as e:
        print(f"描述表错误: {str(e)}")
        return False
    
    finally:
        if cursor:
            cursor.close()
        conn.close()

def truncate_table(table_name, confirm=False):
    """清空表数据"""
    if not confirm:
        response = input(f"确定要清空表 '{table_name}' 吗? 此操作不可撤销! (y/N): ")
        if response.lower() != 'y':
            print("操作已取消")
            return False
    
    conn = get_db_connection()
    if not conn:
        return False
    
    cursor = None
    try:
        cursor = conn.cursor()
        
        # 检查表是否存在
        cursor.execute("""
            SELECT 1 
            FROM information_schema.tables 
            WHERE table_schema = 'public' AND table_name = %s
        """, (table_name,))
        
        if not cursor.fetchone():
            print(f"表 '{table_name}' 不存在")
            return False
        
        # 清空表
        cursor.execute(f"TRUNCATE TABLE {table_name} RESTART IDENTITY CASCADE")
        conn.commit()
        
        print(f"表 '{table_name}' 已清空")
        return True
    
    except Exception as e:
        print(f"清空表错误: {str(e)}")
        conn.rollback()
        return False
    
    finally:
        if cursor:
            cursor.close()
        conn.close()

def export_table_data(table_name, output_file=None):
    """导出表数据到JSON文件"""
    conn = get_db_connection()
    if not conn:
        return False
    
    cursor = None
    try:
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        
        # 检查表是否存在
        cursor.execute("""
            SELECT 1 
            FROM information_schema.tables 
            WHERE table_schema = 'public' AND table_name = %s
        """, (table_name,))
        
        if not cursor.fetchone():
            print(f"表 '{table_name}' 不存在")
            return False
        
        # 获取表数据
        cursor.execute(f"SELECT * FROM {table_name}")
        rows = cursor.fetchall()
        
        # 默认文件名
        if not output_file:
            output_file = f"{table_name}_export.json"
        
        # 导出为JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(rows, f, ensure_ascii=False, indent=2)
        
        print(f"已导出 {len(rows)} 行数据到文件 '{output_file}'")
        return True
    
    except Exception as e:
        print(f"导出表错误: {str(e)}")
        return False
    
    finally:
        if cursor:
            cursor.close()
        conn.close()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="百家乐系统数据库管理工具")
    
    # 子命令
    subparsers = parser.add_subparsers(dest="command", help="子命令")
    
    # 初始化数据库命令
    init_parser = subparsers.add_parser("init", help="初始化数据库表")
    
    # 列出表命令
    list_parser = subparsers.add_parser("list", help="列出所有表")
    
    # 描述表结构命令
    describe_parser = subparsers.add_parser("describe", help="描述表结构")
    describe_parser.add_argument("table", help="表名")
    
    # 清空表命令
    truncate_parser = subparsers.add_parser("truncate", help="清空表数据")
    truncate_parser.add_argument("table", help="表名")
    truncate_parser.add_argument("-y", "--yes", action="store_true", 
                               help="不询问确认直接执行")
    
    # 导出表数据命令
    export_parser = subparsers.add_parser("export", help="导出表数据到JSON文件")
    export_parser.add_argument("table", help="表名")
    export_parser.add_argument("-o", "--output", help="输出文件名")
    
    args = parser.parse_args()
    
    # 处理命令
    if args.command == "init":
        print("初始化数据库...")
        check_database_exists()
        init_database()
    
    elif args.command == "list":
        list_tables()
    
    elif args.command == "describe":
        describe_table(args.table)
    
    elif args.command == "truncate":
        truncate_table(args.table, args.yes)
    
    elif args.command == "export":
        export_table_data(args.table, args.output)
    
    else:
        parser.print_help()

if __name__ == "__main__":
    main() 